#-*- coding: utf-8 -*-
#
# Copyright 2013-2014 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

"""
 This is a 6dof model for a fixed wing vehicle
"""

import math
import numpy as np
import scipy.integrate
import scipy.optimize
import matplotlib.pyplot as plt

import pat.dynamic_model as dm
import pat.frames as fr
import pat.utils as pu
import pat.algebra as pal
import pat.atmosphere as patm

"""
Components of the state vector
"""
sv_x     = 0   # position x axis
sv_y     = 1   # position y axis
sv_z     = 2   # heigh above ground
sv_v     = 3   # airspeed
sv_alpha = 4   # alpha
sv_beta  = 5   # beta
sv_phi   = 6   # roll  (euler, ltp to body)
sv_theta = 7   # pitch (euler, ltp to body)
sv_psi   = 8   # yaw   (euler, ltp to body)
sv_p     = 9   # rotational vel body x
sv_q     = 10  # rotational vel body y
sv_r     = 11  # rotational vel body z
sv_size  = 12

"""
Components of the input vector
"""
iv_da   = 0
iv_de   = 1
iv_dr   = 2
iv_size = 4


def get_aero_to_body(X):
    """
    computes the aero to body rotation matix
    """
    ca = math.cos(X[sv_alpha]); sa = math.sin(X[sv_alpha]) 
    cb = math.cos(X[sv_beta]);  sb = math.sin(X[sv_beta])
    return np.array([[ca*cb, -ca*sb, -sa],
                     [sb   ,  cb   ,  0.],
                     [sa*cb, -sa*sb,  ca]])

def get_f_eng_body(X, U, P):
    """
    return propulsion forces expressed in body frame
    """
    rho = patm.get_rho(-X[sv_z])
    f_engines_body = np.zeros((P.eng_nb, 3))
    for i in range(0, P.eng_nb):
        thrust = U[i]*P.fmaxs[i]*math.pow((rho/P.rhois[i]),P.nrhos[i])*math.pow((X[sv_v]/P.Vis[i]),P.nVs[i]) 
        f_engines_body[i] = np.dot(P.eng_to_body[i], np.array([thrust, 0., 0.]))
    return f_engines_body 

def get_f_aero_body(X, Usfc, P, Pdyn):
    """
    return aerodynamic forces expressed in body frame
    """
    d_alpha = X[sv_alpha] - P.alpha0
    rvel =  X[sv_p:sv_r+1]*np.array([P.Bref, P.Cref, P.Bref])/2/P.Vref

    CL = P.CL0 + P.CL_alpha*d_alpha + P.CL_beta*X[sv_beta] +\
         np.dot(P.CL_omega,rvel) + np.dot(P.CL_sfc,Usfc)

    CD = P.CD0 + P.CD_k1*CL + P.CD_k2*(CL**2) + np.dot(P.CD_sfc,Usfc)

    CY = P.CY_alpha*d_alpha + P.CY_beta*X[sv_beta] +\
         np.dot(P.CY_omega,rvel) + np.dot(P.CY_sfc,Usfc)
    
    return Pdyn*P.Sref*np.dot(get_aero_to_body(X),[-CD, CY, -CL])

def get_m_eng_body(f_eng_body, P):
    """
    return propulsion moments expressed in body frame
    """
    m = np.zeros(3)
    for i in range(0, P.eng_nb):
        m += np.cross(P.eng_pos[i], f_eng_body[i])
    return m

def get_m_aero_body(X, Usfc, P, Pdyn):
    """
    return aerodynamic moments expressed in body frame
    """
    d_alpha = X[sv_alpha] - P.alpha0
    rvel =  X[sv_p:sv_r+1]*np.array([P.Bref, P.Cref, P.Bref])/2/P.Vref

    Cl =         P.Cl_alpha*d_alpha + P.Cl_beta*X[sv_beta] +\
         np.dot(P.Cl_omega,rvel) + np.dot(P.Cl_sfc,Usfc)
    Cm = P.Cm0 + P.Cm_alpha*d_alpha + P.Cm_beta*X[sv_beta] +\
         np.dot(P.Cm_omega,rvel) + np.dot(P.Cm_sfc,Usfc)
    Cn =         P.Cn_alpha*d_alpha + P.Cn_beta*X[sv_beta] +\
         np.dot(P.Cn_omega,rvel) + np.dot(P.Cn_sfc,Usfc)

    return Pdyn*P.Sref*np.array([Cl*P.Bref, Cm*P.Cref, Cn*P.Bref])

def dyn(X, t, U, P):
    """
    Dynamic model
    """
    rho = patm.get_rho(-X[sv_z])
    Pdyn = 0.5*rho*X[sv_v]**2

    Ueng = U[0:P.eng_nb]                    # engines part of input vector
    Usfc = U[P.eng_nb:P.eng_nb+P.sfc_nb]    # control surfaces part of input vector
    X_rvel_body = X[sv_p:sv_r+1]            # body rotational velocities
    X_euler = X[sv_phi:sv_psi+1]            # euler angles                          

    # Newton for forces in body frame
    f_aero_body = get_f_aero_body(X, Usfc, P, Pdyn)
    f_eng_body  = get_f_eng_body(X, Ueng, P)
    earth_to_body = pal.rmat_of_euler(X_euler)
    f_weight_body = np.dot(earth_to_body, [0., 0., P.m*P.g])
    forces_body = f_aero_body + np.sum(f_eng_body, axis=0) + f_weight_body
    
    vel_body = np.dot(get_aero_to_body(X), [X[sv_v], 0., 0.])  # u, v, w
    accel_body = 1./P.m*forces_body - np.cross(X_rvel_body, vel_body)

    # Newton for moments in body frame
    m_aero_body = get_m_aero_body(X, Usfc, P, Pdyn)
    m_eng_body = get_m_eng_body(f_eng_body, P)
    raccel_body = np.dot(P.invI, m_aero_body + m_eng_body - np.cross(X_rvel_body, np.dot(P.I, X_rvel_body)))
    
    Xdot = np.zeros(sv_size)
    Xdot[sv_x:sv_z+1] = np.dot(np.transpose(earth_to_body), vel_body)

    Xdot[sv_v] = np.inner(vel_body, accel_body)/X[sv_v]
    u, v, w = vel_body
    ud, vd, wd = accel_body
    Xdot[sv_alpha] = (u*wd - w*ud)/(u**2+w**2)
    Xdot[sv_beta] = (X[sv_v]*vd - v*Xdot[sv_v]) / X[sv_v] / math.sqrt(u**2+w**2)

    Xdot[sv_phi:sv_psi+1] = pal.euler_derivatives(X_euler, X_rvel_body)
    
    Xdot[sv_p:sv_r+1] = raccel_body

    return Xdot


def trim(P, args=None, debug=False):
    """
    Find throttle, elevator  and angle of attack corresponding
    to the given airspeed and and flight path
    """
    if args<>None:
        va, gamma, h = (args['va'], args['gamma'], args['h'] )
    else:
        va, gamma, h = (P.Vref, 0., 0.)

    if debug:
        print "searching for constant path trajectory with"
        print "  va     {:f} m/s".format(va)
        print "  gamma  {:f} deg".format(pu.deg_of_rad(gamma))

    def err_func((throttle, elevator, alpha)):
        X=[0., 0., -h, va, alpha, 0., 0.,  gamma+alpha, 0., 0., 0., 0.]
        U = np.zeros(P.input_nb)
        U[0:P.eng_nb] = throttle; U[P.eng_nb+iv_de] = elevator
        Xdot = dyn(X, 0., U, P)
        Xdot_ref = [va*math.cos(gamma), 0., -va*math.sin(gamma), 0., 0., 0., 0., 0., 0., 0., 0., 0.]
        return np.linalg.norm(Xdot - Xdot_ref)

    p0 = [0.2, pu.rad_of_deg(2.), pu.rad_of_deg(0.)]
    thr_e, ele_e, alpha_e = scipy.optimize.fmin_powell(err_func, p0, disp=debug, ftol=1e-9)

    if debug:
        print """result:
  throttle        : {:f} %
  elevator        : {:f} deg
  angle of attack : {:f} deg""".format(100.*thr_e, pu.deg_of_rad(ele_e), pu.deg_of_rad(alpha_e))

    Ue = np.zeros(P.input_nb)
    Ue[0:P.eng_nb] = thr_e; Ue[P.eng_nb+iv_de] = ele_e
    Xe = [va*math.cos(gamma), 0., va*math.sin(gamma), va, alpha_e, 0., 0., gamma+alpha_e, 0., 0., 0., 0.]  
    return Xe, Ue


import pat.vehicles.fixed_wing.dynamic_model_python_parameters
class Param(pat.vehicles.fixed_wing.dynamic_model_python_parameters.Param):
    pass




class DynamicModel(dm.DynamicModel):

    sv_x     = sv_x     # position x axis
    sv_y     = sv_y     # position y axis
    sv_z     = sv_z     # heigh above ground
    sv_v     = sv_v     # airspeed
    sv_alpha = sv_alpha # alpha
    sv_beta  = sv_beta  # beta
    sv_phi   = sv_phi   # roll  (euler, ltp to body)
    sv_theta = sv_theta # pitch (euler, ltp to body)
    sv_psi   = sv_psi   # yaw   (euler, ltp to body)
    sv_p     = sv_p     # rotational vel body x
    sv_q     = sv_q     # rotational vel body y
    sv_r     = sv_r     # rotational vel body z
    sv_size  = sv_size

    
    iv_th   = 0    # throttle
    iv_da   = 1    # aileron
    iv_de   = 2    # elevator
    iv_dr   = 3    # rudder
    iv_size = 4
    # hack for multiple engines
    _iv_da = 0
    _iv_de = 1
    _iv_dr = 2

    dyn = lambda self, X, t, U, P: dyn(X, t, U, self.P)
    trim = lambda self, args=None, debug=False: trim(self.P, args, debug)

    def __init__(self, params=None):
        print "Info: Dynamic fixed wing basic"
        dm.DynamicModel.__init__(self)
        if params == None: params="../config/Rcam_single_engine.xml"
        self.X = np.zeros(DynamicModel.sv_size)
        self.P = Param(params)
        self.reset()

    def name(self):
        return "Fixed Wing Python Basic ({:s})".format(self.P.name)

    def reset(self, X0=None):
        if X0<>None: self.X = X0
        else: self.X = np.array([0., 0., 0., 68., 0., 0., 0., 0., 0., 0., 0., 0.])
        return self.X

    def run(self, dt, U):
        foo, self.X = scipy.integrate.odeint(dyn, self.X, [0, dt], args=(U, self.P, ))
        return self.X

    def param(self):
        return str(self.P)

    def iv_dth(self):
        if self.P.eng_nb>1: return range(0,self.P.eng_nb)
        else: return 0
    def iv_da(self): return self.P.eng_nb + DynamicModel._iv_da
    def iv_de(self): return self.P.eng_nb + DynamicModel._iv_de
    def iv_dr(self): return self.P.eng_nb + DynamicModel._iv_dr
    def input_nb(self): return self.P.input_nb

    def state_SixDOFfEuclidianEuler(self):
        X = np.zeros(fr.SixDOFfEuclidianEuler.size)
        X[fr.SixDOFfEuclidianEuler.x:fr.SixDOFfEuclidianEuler.z+1] = self.X[sv_x:sv_z+1]
        X[fr.SixDOFfEuclidianEuler.phi:fr.SixDOFfEuclidianEuler.r+1] = self.X[sv_phi:sv_r+1]
        return X

    def get_jacobian(self, Xe, Ue):
        A,B = pu.num_jacobian(Xe, Ue, self.P, dyn)
        return A, B 

    def state_str(self):
        return """pos: {:-.2f}, {:-.2f}, {:-.2f} m
        vel: {:-.2f} m/s, alpha {:-.2f}, beta {:-.2f} deg
        att:    {:-.2f}, {:-.2f}, {:-.2f} deg
        """.format(self.X[sv_x], self.X[sv_y], self.X[sv_z],
                   self.X[sv_v], pu.deg_of_rad(self.X[sv_alpha]), pu.deg_of_rad(self.X[sv_beta]),
                   pu.deg_of_rad(self.X[sv_phi]), pu.deg_of_rad(self.X[sv_theta]), pu.deg_of_rad(self.X[sv_psi]))

    def plot_trajectory(self, time, X, U=None, figure=None, window_title="Trajectory", legend=None, filename=None): 
        plot_trajectory(time, X, U, figure, window_title, legend, filename)

#
# Some plotting functions
#
def plot_trajectory(time, X, U=None, figure=None, window_title="Trajectory",
                    legend=None, filename=None):

        margins=(0.04, 0.05, 0.98, 0.96, 0.20, 0.34)
        figure = pu.prepare_fig(figure, window_title, figsize=(20.48, 10.24), margins=margins)

        plots = [("x", "m", X[:,sv_x]), ("y", "m", X[:,sv_y]), ("z", "m", X[:,sv_z]),
                 ("v",         "m/s",               X[:,sv_v]),
                 ("$\\alpha$", "deg", pu.deg_of_rad(X[:,sv_alpha])),
                 ("$\\beta$",  "deg", pu.deg_of_rad(X[:,sv_beta])),
                 ("$\phi$",    "deg", pu.deg_of_rad(X[:,sv_phi])),
                 ("$\\theta$", "deg", pu.deg_of_rad(X[:,sv_theta])),
                 ("$\\psi$",   "deg", pu.deg_of_rad(X[:,sv_psi])),
                 ("$p$",     "deg/s", pu.deg_of_rad(X[:,sv_p])),
                 ("$q$",     "deg/s", pu.deg_of_rad(X[:,sv_q])),
                 ("$r$",     "deg/s", pu.deg_of_rad(X[:,sv_r]))]

        nrow = 5 if U<>None else 4
        for i, (title, ylab, data) in enumerate(plots):
            ax = plt.subplot(nrow, 3, i+1)
            plt.plot(time, data)
            pu.decorate(ax, title=title, ylab=ylab)
            
        if legend<>None:
            plt.legend(legend, loc='best')

        if U<>None:
            ax = figure.add_subplot(5, 3, 13)
            ax.plot(time, 100*U[:, 0])
            pu.decorate(ax, title="$d_{th}$", ylab="%")
            ax = figure.add_subplot(5, 3, 14)
            ax.plot(time, pu.deg_of_rad(U[:, iv_da+1]))
            pu.decorate(ax, title="$d_a$", ylab="deg")
            ax = figure.add_subplot(5, 3, 15)
            ax.plot(time, pu.deg_of_rad(U[:, iv_de+1]))
            pu.decorate(ax, title="$d_e$", ylab="deg")
    
        return figure
        


